"""Wrapper around PaddleNLP embedding models."""
from typing import Any, List
from langchain.embeddings.base import Embeddings
from pydantic import BaseModel, Extra

class PaddleNLPEmbeddings(BaseModel, Embeddings):
    """Wrapper around paddlenlp embedding models."""
    text_encoder: Any
    model: str ='rocketqa-zh-base-query-encoder'
    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        try:
            import paddle.nn.functional as F
            from paddlenlp import Taskflow
            self.text_encoder = Taskflow("feature_extraction", model=self.model)
        except ImportError as e:
            raise ValueError(
                "Could not import some python packages." "Please install it with `pip install paddlenlp`."
            ) from e
    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        texts = list(map(lambda x: x.replace("\n", " "), texts))
        text_embeds = self.text_encoder(texts)
        embeddings = text_embeds["features"].numpy()
        return embeddings.tolist()
    def embed_query(self, text: str) -> List[float]:
        text = text.replace("\n", " ")
        text_embeds = self.text_encoder(text)
        embedding = text_embeds["features"].numpy()[0]
        return embedding.tolist()
		
		